#coding:utf-8
'''
 获取数据 s&p500
 对 s&p500 数据进行处理，形成特征
'''
import tushare
import pandas as pd
import matplotlib.pyplot as plt

data_dir = './Data/sp500/'
file = data_dir + 'SP500.csv'
file_processed = data_dir + 'SP500_processed.csv'
ratio = 0.2
'''
    <class 'pandas.core.frame.DataFrame'>
    RangeIndex: 16862 entries, 0 to 16861
    Data columns (total 7 columns):
    Date         16862 non-null object
    Open         16862 non-null float64
    High         16862 non-null float64
    Low          16862 non-null float64
    Close        16862 non-null float64
    Adj Close    16862 non-null float64
    Volume       16862 non-null int64
    dtypes: float64(5), int64(1), object(1)
    memory usage: 922.2+ KB
    
    Structure
           Date       Open       High        Low      Close  Adj Close   Volume
    0  1950-12-18  19.850000  19.850000  19.850000  19.850000  19.850000  4500000
    1  1950-12-19  19.959999  19.959999  19.959999  19.959999  19.959999  3650000
    2  1950-12-20  19.969999  19.969999  19.969999  19.969999  19.969999  3510000
'''
def raw_data():
    data_raw = pd.read_csv(file)
    # print(data_raw.info())
    # 因为数据是从1950年开始的，需要对数据进行反转
    data_raw = data_raw.sort_index(ascending=False, axis=0)
    # 日期作为索引值
    # data_raw.index = data_raw['Date']
    # 存储数据
    # data_raw.to_csv(file_processed)
    return data_raw

'''
    :arg
    n 代表将加入前n天的特征
'''
def before_n_days(n, data):
    new_data = []
    temp = data[0:-3]
    data = data['Close'].values
    data = list(data)
    for i in range(len(data[n:])):
        new_data_item = []
        new_data_item = data[i+1:i+n]
        label = 1 if data[i] - data[i+1]>0 else 0
        new_data_item.extend([label])
        new_data.append(new_data_item)
    new_data = pd.DataFrame(new_data, columns=['Close-1', 'Close-2', 'label'])
    new_data.index = temp.index
    return new_data

# data_raw = raw_data()
# data_raw = data_raw.drop(columns=['Date', 'Volume', 'Adj Close'])
# before_n_days(3, data_raw)

# 过去的数据作为训练集合，最近的数据作为测试集合
def split_data(ratio=0.2):
    data_raw = raw_data()
    data_raw = data_raw.drop(columns=['Date', 'Volume', 'Adj Close'])

    # 3 代表提前几天数据
    close_label = before_n_days(3, data_raw)
    data_raw = pd.concat([data_raw, close_label], axis=1)
    # 增加特征
    data_raw['outcome'] = data_raw['Open'] - data_raw['Close']
    data_raw['gain'] = data_raw['outcome']
    data_raw.loc[:, 'outcome'][data_raw.loc[:, 'outcome'] >= 0] = 1
    data_raw.loc[:, 'outcome'][data_raw.loc[:, 'outcome'] < 0] = 0

    data_raw['HO'] = data_raw['High'] - data_raw['Open']
    data_raw['LO'] = data_raw['Low'] - data_raw['Open']

    # 错位相减获得label
    # data_raw_1 = data_raw.shift(1)
    data_raw_2 = data_raw.shift(2)

    # data_raw['feat1'] = data_raw['Close'] >= data_raw_1['Close']
    data_raw['label2'] = data_raw['Close'] >= data_raw_2['Close']

    data_raw = data_raw.sort_index(ascending=False, axis=0)
    # 测试数据和训练数据的划分，还是不用划分
    train_data = data_raw[round(len(data_raw)*ratio):]
    test_data = data_raw[1:round(len(data_raw)*ratio)]

    return train_data, test_data, data_raw

# from STOCK_SVM_tushare
def refact_data(Data):
	#Create more data
	value = pd.Series(Data['close'].shift(-1)-Data['close'],index=Data.index)	# Series用于数据平移, 错位相减
	#Data['Next_Open'] = Data['Open'].shift(-1) #Next day's Open data.
	Data['high-low'] = Data['high']-Data['low'] #Difference between High and Low
	Data['NOpen-Close']=Data['open'].shift(-1)-Data['close'] #Next Day's Open-today's Close
	Data['close-YClose']=Data['close']-Data['close'].shift(1) #Today is rise or fall
	Data['close-open']=Data['close']-Data['open'] #today's Close - Open
	Data['High-close'] = Data['high']-Data['close'] #today's High - Close
	Data['close-Low'] = Data['close']-Data['low'] #today's Close - Low
	value[value>=0]=1 #0 means rise
	value[value<0]=0 #1 means fall
	Data=Data.dropna(how='any') 	# 删除缺失数据
	del(Data['open'])
	del(Data['close'])
	del(Data['high'])
	del(Data['low'])
	#print(Data)
	print(type(Data))
	return Data, value

# def hello():
#     print("utils / util / hello")

# split_data(0.2)